package ua.mykola.graph;

import java.util.HashSet;
import java.util.PriorityQueue;
import java.util.Set;

public class PrimMST {

	private Graph graph;
	private Vertex startVertex;
	private PriorityQueue<Edge> queue;
	private Set<Edge> mstEdges = new HashSet<Edge>();
	private Set<Vertex> mstVertices = new HashSet<Vertex>();

	public PrimMST(Graph graph, String label) {
		this.graph = graph;
		startVertex = new Vertex(label);
		queue = new PriorityQueue<Edge>(11, new EdgeWeightComparator());
	}

	public void traverse() {
		mstVertices.add(startVertex);
		queue.addAll(graph.getEdgesFrom(startVertex));

		while (!queue.isEmpty()) {
			Edge edge = queue.poll();
			Vertex vertexA = edge.getVertexA();
			Vertex vertexB = edge.getVertexB();
			if (mstVertices.contains(vertexA) && mstVertices.contains(vertexB)) {
				continue;
			}
			mstEdges.add(edge);
			if (mstVertices.contains(vertexA)) {
				mstVertices.add(vertexB);
				queue.addAll(graph.getEdgesFrom(vertexB));
			} else if (mstVertices.contains(vertexB)) {
				mstVertices.add(vertexA);
				queue.addAll(graph.getEdgesFrom(vertexA));
			} else {
				throw new IllegalStateException();
			}
		}

	}

	public Set<Edge> getMSTEdges() {
		return mstEdges;
	}

	public int getMSTWeight() {
		int weight = 0;
		for (Edge edge : mstEdges) {
			weight += edge.getWeight();
		}
		return weight;
	}

}
